package main

import (
    "io"
    "flag"
    "fmt"
    "log"
    "path"
    "strings"
    "os"
    "context"
    "os/exec"
    "net/url"
    "net/http"
    "path/filepath"
    "github.com/gorilla/mux"
)

var (
    PORT string
    UP_STREAM string
)
func init() {
    // 在init函数中初始化命令行参数
    flag.StringVar(&PORT, "p", "8080", "The port to listen on")
}

func removeUpstreamQueryParam(urlStr string) string {
	upstreamIndex := strings.Index(urlStr, "upstream=")
	ampIndex := strings.Index(urlStr[upstreamIndex:], "&")
	if ampIndex == -1 {
            // 如果没有找到 &，说明 upstream 是最后一个参数，直接去掉它和&
            return urlStr[:upstreamIndex - 1]
	}
        return urlStr
}

func streamHandler(w http.ResponseWriter, r *http.Request) {
    //vars := mux.Vars(r)
    //upstream := vars["upstream"]
    upstream := r.URL.Query().Get("upstream")
    if upstream == "" {
        http.Error(w, "upstream parameter is missing", http.StatusBadRequest)
        return
    }

    fullURL := r.URL.String()
    fullURL = removeUpstreamQueryParam(fullURL)
    fmt.Println(fullURL)

    w.Header().Set("Access-Control-Allow-Origin", "*")
    w.Header().Set("Access-Control-Allow-Methods", "POST, GET, PUT, OPTIONS")
    w.Header().Set("Access-Control-Allow-Headers", "Keep-Alive,User-Agent,If-Modified-Since,Cache-Control,Content-Type,Authorization")
    w.Header().Set("Access-Control-Expose-Headers", "Content-Length, Access-Control-Allow-Origin, Access-Control-Allow-Headers, Content-Type")
    w.Header().Set("Access-Control-Allow-Credentials", "true")

    w.Header().Set("Transfer-Encoding", "chunked")
    w.Header().Set("Content-Type", "video/x-flv") // 根据流的类型设置适当的Content-Type
    w.WriteHeader(http.StatusOK)

    flusher, ok := w.(http.Flusher)
    if !ok {
        http.Error(w, "Streaming unsupported!", http.StatusInternalServerError)
        return
    }
    flusher.Flush()
    UP_STREAM = "http://" + upstream
    up, err := url.Parse(UP_STREAM)
    if err != nil {
        fmt.Println(err)
        return
    }
    up.Path = path.Join(up.Path, fullURL)
    i := up.String()
    inputStreamURL, err:=  url.QueryUnescape(i)

    ctx, cancel:= context.WithCancel(context.Background())
    defer cancel() // 确保在函数返回时取消context

    cwd, err := os.Getwd()
    if err != nil {
        panic(err)
    }
    ffmpegPath := filepath.Join(cwd, "bin", "bin", "ffmpeg")
    // 创建ffmpeg命令，将直播流捕获到标准输出
    fmt.Println(ffmpegPath, inputStreamURL)
    cmd := exec.CommandContext(ctx, ffmpegPath, "-i", inputStreamURL, "-c:v", "libx264", "-preset", "veryfast", "-g", "25", "-bf", "0", "-an", "-f", "flv", "-")

    // 创建一个管道来读取ffmpeg的标准输出
    pr, pw, err := os.Pipe()
    if err != nil {
        http.Error(w, "Error creating pipe for ffmpeg output", http.StatusInternalServerError)
        return
    }
    defer pr.Close()

    // 将ffmpeg的标准输出重定向到管道
    cmd.Stdout = pw

    // 启动ffmpeg命令
    if err := cmd.Start(); err != nil {
        http.Error(w, "Error starting ffmpeg", http.StatusInternalServerError)
        return
    }
    defer cmd.Wait() // 确保等待ffmpeg命令完成
    defer pw.Close() // 在cmd.Wait()之后关闭写端，确保数据被完全写入管道


    go func() {
        _, err := io.Copy(w, pr)
	if err != nil {
	    log.Printf("Error forwarding stream: %v", err)
	    cancel() // 当io.Copy返回错误时，取消context以停止ffmpeg
	}
    }()

	// 监听客户端断开连接

    notify := w.(http.CloseNotifier).CloseNotify()
    <-notify // 等待客户端断开连接
    cancel() // 取消context以停止ffmpeg
}

func isURLValid(toTest string) bool {
    _, err := url.ParseRequestURI(toTest)
    if err != nil {
        return false
    }
    u, err := url.Parse(toTest)
    if err != nil {
        return false
    }
    return u.Scheme == "http" || u.Scheme == "https"
}

func main() {
    flag.Parse()
    if PORT == "" {
        fmt.Fprintln(os.Stderr, "Error: -p or --p is required\nExample:\n./forwarder  -p 9090 ")
        flag.PrintDefaults() // 打印默认参数值  
        os.Exit(1)
    }

    r := mux.NewRouter()
    r.PathPrefix("/").HandlerFunc(streamHandler)

    log.Println("Server started on :" + PORT)
    if err := http.ListenAndServe(":" + PORT, r); err != nil {
        log.Fatal(err)
    }
}
